"""

"""
import torch
import tiktoken
from torch.utils.data import Dataset, DataLoader

vocab_size = 50257
output_dim = 256
token_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)

max_length = 4


class GPTDatasetV2(Dataset):
    def __init__(self, txt, tokenizer, max_length, stride):
        self.input_ids = []
        self.target_ids = []

        # tokenize the entire text
        token_ids = tokenizer.encode(txt)

        # use a sliding window to chunk the book into overlapping sequences of max_length
        for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i: i+max_length]
            target_chunk = token_ids[i+1: i+max_length+1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

    def __len__(self):
        """
        return total number of rows in the dataset
        """
        return len(self.input_ids)

    def __getitem__(self, idx):
        """
        return a single row from the dataset
        """
        return self.input_ids[idx], self.target_ids[idx]


def create_dataloader_v2(txt, batch_size=4, max_length=256, stride=128, shuffle=True, drop_last=True, num_workers=0):
    # initialize the tokenizer, it's vocabulary that map token to id
    tokenizer = tiktoken.get_encoding("gpt2")

    # create dataset
    dataset = GPTDatasetV2(txt, tokenizer, max_length, stride)

    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=0)


def TestVerdict2():
    # raw text
    with open("the-verdict.txt", 'r', encoding='utf-8') as f:
        raw_text = f.read()

    max_length = 4

    dataloader = create_dataloader_v2(raw_text, batch_size=8, max_length=max_length, stride=max_length, shuffle=False)
    data_iter = iter(dataloader)
    inputs, targets = next(data_iter)
    print("Inputs:\n", inputs)
    print("\nTargets:\n", targets)
    print(("\nInputs shape:\n", inputs.shape))

    token_embeddings = token_embedding_layer(inputs)
    print(token_embeddings.shape)

    context_length = max_length
    pos_embedding_layer = torch.nn.Embedding(context_length, output_dim)
    # torch.arange(context_length) : [0, 1, 2, ..., context_length-1]
    pos_embeddings = pos_embedding_layer(torch.arange(context_length))
    print(pos_embeddings.shape)

    input_embeddings = token_embeddings + pos_embeddings
    print(input_embeddings.shape)

TestVerdict2()